{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pysam\n",
    "import numpy as np\n",
    "from deepsignal3.utils.process_utils import complement_seq\n",
    "from deepsignal3.utils.ref_reader import get_contig2len,get_contig2len_n_seq\n",
    "from deepsignal3.utils.process_utils import get_motif_seqs\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "reference_path='/homeb/hpc/refs/genome/chm13v2.0.fa'\n",
    "chrom2len,contigs = get_contig2len_n_seq(reference_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'TT'"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "contigs['chr11'][86669225 -1:86669225+1 ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'AT'"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "contigs['chr11'][86669219 -1:86669219+1 ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'GG'"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "contigs['chr11'][86669201 -1:86669201+1 ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'CG'"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "contigs['chr11'][86670261 -1:86670261+1 ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "bamfile = pysam.AlignmentFile('/homeb/hpc/raw_data/hg002_r1041_4khz/demo/100/100.demo.dorado.5mc_5hmc.sup.bam', \"rb\", check_sq=False)\n",
    "bam_index=pysam.IndexedReads(bamfile)\n",
    "bam_index.build()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_refloc_of_methysite_in_motif(seqstr, motifset, methyloc_in_motif=0):\n",
    "    \"\"\"\n",
    "\n",
    "    :param seqstr:\n",
    "    :param motifset:\n",
    "    :param methyloc_in_motif: 0-based\n",
    "    :return:\n",
    "    \"\"\"\n",
    "    motifset = set(motifset)\n",
    "    strlen = len(seqstr)\n",
    "    motiflen = len(list(motifset)[0])\n",
    "    sites = []\n",
    "    for i in range(0, strlen - motiflen + 1):\n",
    "        if seqstr[i : i + motiflen] in motifset:\n",
    "            sites.append(i + methyloc_in_motif)\n",
    "    return sites\n",
    "\n",
    "def align_signals_and_extend_ref_seq(pos_pair, read_seq, ref_seq,motif_seqs,methyloc,strand,ref_start,ref_end):\n",
    "    # 确定最前端和最末端有效的比对索引\n",
    "    first_valid_index = next((i for i, (_, ref_pos) in enumerate(pos_pair) if ref_pos is not None), len(pos_pair))\n",
    "    last_valid_index = len(pos_pair) - 1 - next((i for i, (_, ref_pos) in enumerate(reversed(pos_pair)) if ref_pos is not None), len(pos_pair))\n",
    "\n",
    "    # 生成新的 ref_seq 和 ref_signal\n",
    "    new_ref_seq = []\n",
    "    new_ref_signal = []\n",
    "\n",
    "    # 填充前端未比对的部分\n",
    "    for i in range(first_valid_index):\n",
    "        read_pos, _ = pos_pair[i]\n",
    "        if read_pos is not None:\n",
    "            new_ref_seq.append(read_seq[read_pos])\n",
    "            \n",
    "\n",
    "    last_valid_ref_pos = len(new_ref_seq) - 1\n",
    "\n",
    "    # 处理中间的比对部分\n",
    "    for i in range(first_valid_index, last_valid_index + 1):\n",
    "        read_pos, ref_pos = pos_pair[i]\n",
    "\n",
    "        if ref_pos is not None:\n",
    "            # 确保 new_ref_seq 的长度足够\n",
    "            # while len(new_ref_seq) < ref_pos:\n",
    "            #     new_ref_seq.append(None)  # 用 None 占位，表示插入的部分\n",
    "            #     new_ref_signal.append([]) # 对应的信号也填充为空列表\n",
    "            \n",
    "            new_ref_seq.append(ref_seq[ref_pos])\n",
    "            \n",
    "            last_valid_ref_pos = len(new_ref_seq) - 1\n",
    "\n",
    "\n",
    "\n",
    "    # 填充后端未比对的部分\n",
    "    for i in range(last_valid_index + 1, len(pos_pair)):\n",
    "        read_pos, _ = pos_pair[i]\n",
    "        if read_pos is not None:\n",
    "            new_ref_seq.append(read_seq[read_pos])\n",
    "\n",
    "            \n",
    "\n",
    "    new_ref_seq = ''.join(base  for base in new_ref_seq)\n",
    "    ref_readlocs = dict()\n",
    "    ref_poss = []\n",
    "    pred_pos = []\n",
    "    ref_pos = -1\n",
    "    tsite_locs = get_refloc_of_methysite_in_motif(\n",
    "        new_ref_seq, set(motif_seqs), methyloc)\n",
    "    for loc_in_read in tsite_locs:\n",
    "        if loc_in_read<first_valid_index:\n",
    "            ref_pos = -1\n",
    "            ref_poss.append(ref_pos)\n",
    "            pred_pos.append(loc_in_read)\n",
    "            continue\n",
    "        if loc_in_read>last_valid_index:\n",
    "            ref_pos = -1\n",
    "            ref_poss.append(ref_pos)\n",
    "            pred_pos.append(loc_in_read)\n",
    "            continue\n",
    "        if strand == \"-\":\n",
    "            ref_pos = ref_end-loc_in_read-1+first_valid_index\n",
    "        else:\n",
    "            ref_pos = ref_start+loc_in_read-first_valid_index\n",
    "        #ref_readlocs[loc_in_read+first_valid_index] = ref_pos\n",
    "        ref_poss.append(ref_pos)\n",
    "        pred_pos.append(loc_in_read)\n",
    "    #new_ref_seq, pred_pos,n_positions=filter_n_and_update_indices(new_ref_seq, pred_pos)\n",
    "    ref_readlocs = dict(zip(pred_pos, ref_poss))\n",
    "    #new_ref_signal = [new_ref_signal[i] for i in range(len(new_ref_signal)) if i not in n_positions]\n",
    "    #n_lens=len(n_positions)\n",
    "    return new_ref_seq, new_ref_signal, ref_readlocs, ref_poss, pred_pos#, n_lens"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "read_iter=bam_index.find('a70efe35-e524-410f-8472-1bb8b9edac3b')\n",
    "motif_seqs = get_motif_seqs('C', True)\n",
    "for bam_read in read_iter:\n",
    "\n",
    "    ref_readlocs=dict()\n",
    "\n",
    "    if bam_read.is_unmapped:\n",
    "        continue\n",
    "    reference_name=bam_read.reference_name\n",
    "    strand_code = -1 if bam_read.is_reverse else 1\n",
    "                \n",
    "    strand = \"-\" if bam_read.is_reverse else \"+\"\n",
    "    seq=bam_read.get_forward_sequence() #if  bam_read.get_forward_sequence() is not None else seq      \n",
    "    q_seq=bam_read.query_sequence #if bam_read.query_sequence is not None else q_seq\n",
    "    if bam_read.is_reverse:\n",
    "        ref_seq = complement_seq(bam_read.get_reference_sequence().upper())\n",
    "    else:\n",
    "        ref_seq = bam_read.get_reference_sequence().upper()\n",
    "    if bam_read.has_tag('MD'):\n",
    "        rseq=bam_read.get_reference_sequence()\n",
    "    else:\n",
    "        rseq=''\n",
    "    if rseq is not None and rseq!='':\n",
    "        rseq_complement=complement_seq(rseq)\n",
    "    q_seq_complement=complement_seq(q_seq)\n",
    "    #find_key=(read_name,reference_name)\n",
    "    ref_start=bam_read.reference_start\n",
    "    ref_end = bam_read.reference_end\n",
    "    cigar_tuples = bam_read.cigartuples\n",
    "    qalign_start = bam_read.query_alignment_start\n",
    "    qalign_end = bam_read.query_alignment_end\n",
    "    if bam_read.is_reverse:\n",
    "        seq_start = len(seq) - qalign_end\n",
    "        seq_end = len(seq) - qalign_start\n",
    "    else:\n",
    "        seq_start = qalign_start\n",
    "        seq_end = qalign_end\n",
    "    pos_pair=[]\n",
    "    for read_pos,ref_pos in bam_read.get_aligned_pairs():\n",
    "        # if ref_pos is None:\n",
    "        #     continue\n",
    "        if read_pos is None:\n",
    "            if bam_read.is_reverse:\n",
    "                pos_pair.append((None,ref_end-ref_pos-1))\n",
    "            else:\n",
    "                pos_pair.append((None,ref_pos-ref_start))\n",
    "            continue\n",
    "        if ref_pos is None:\n",
    "            if bam_read.is_reverse:\n",
    "                pos_pair.append((len(seq)-read_pos-1,None))\n",
    "            else:\n",
    "                pos_pair.append((read_pos,None))\n",
    "            continue\n",
    "        if bam_read.is_reverse:\n",
    "            pos_pair.append((len(seq)-read_pos-1,ref_end-ref_pos-1))\n",
    "        else:\n",
    "            pos_pair.append((read_pos,ref_pos-ref_start))\n",
    "    #pos_pair=bam_read.get_aligned_pairs()\n",
    "    #signal_group=align_signals(pos_pair,signal_group,ref_seq)\n",
    "    if strand == \"-\":\n",
    "        pos_pair.reverse()\n",
    "    seq,signal_group, ref_readlocs, ref_poss, pred_pos=align_signals_and_extend_ref_seq(pos_pair, seq, ref_seq,motif_seqs,0,strand,ref_start,ref_end)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[126657298, 126657295, 126657282, 126657281, 126657265, 126657260, 126657257, 126657254, 126657247, 126657238, 126657236, 126657233, 126657220, 126657219, 126657206, 126657203, 126657202, 126657192, 126657190, 126657186, 126657185, 126657157, 126657154, 126657139, 126657138, 126657131, 126657129, 126657128, 126657127, 126657126, 126657125, 126657120, 126657114, 126657112, 126657103, 126657102, 126657091, 126657084, 126657071, 126657068, 126657063, 126657062, 126657039, 126657035, 126657034, 126657030, 126657012, 126657011, 126657002, 126656997, 126656989, 126656988, 126656977, 126656970]\n"
     ]
    }
   ],
   "source": [
    "print(ref_poss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "AG\n",
      "\n",
      "AG\n",
      "\n",
      "GG\n",
      "\n",
      "TG\n",
      "\n",
      "TG\n",
      "\n",
      "CG\n",
      "\n",
      "TG\n",
      "\n",
      "CG\n",
      "\n",
      "AG\n",
      "\n",
      "AG\n",
      "\n",
      "CG\n",
      "\n",
      "AG\n",
      "\n",
      "GG\n",
      "\n",
      "TG\n",
      "\n",
      "AG\n",
      "\n",
      "GG\n",
      "\n",
      "AG\n",
      "\n",
      "AG\n",
      "\n",
      "TG\n",
      "\n",
      "GG\n",
      "\n",
      "AG\n",
      "\n",
      "AG\n",
      "\n",
      "TG\n",
      "\n",
      "GG\n",
      "\n",
      "AG\n",
      "\n",
      "AG\n",
      "\n",
      "GG\n",
      "\n",
      "GG\n",
      "\n",
      "GG\n",
      "\n",
      "GG\n",
      "\n",
      "TG\n",
      "\n",
      "AG\n",
      "\n",
      "AG\n",
      "\n",
      "AG\n",
      "\n",
      "GG\n",
      "\n",
      "TG\n",
      "\n",
      "AG\n",
      "\n",
      "AG\n",
      "\n",
      "AG\n",
      "\n",
      "AG\n",
      "\n",
      "GG\n",
      "\n",
      "TG\n",
      "\n",
      "TG\n",
      "\n",
      "GG\n",
      "\n",
      "AG\n",
      "\n",
      "CG\n",
      "\n",
      "GG\n",
      "\n",
      "TG\n",
      "\n",
      "AG\n",
      "\n",
      "TG\n",
      "\n",
      "GG\n",
      "\n",
      "AG\n",
      "\n",
      "AG\n",
      "\n",
      "AG\n",
      "\n"
     ]
    }
   ],
   "source": [
    "for i in ref_poss:\n",
    "    print(contigs[reference_name][i-1:i+1],end='')\n",
    "    print('\\n')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "deepsignal3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
